import numpy as np
import numba
from collections import deque
from typing import Dict, List, Tuple
import time
from dataclasses import dataclass
from numba import jit
from concurrent.futures import ThreadPoolExecutor


@dataclass
class PlatParaInfo:
    """我方平台信息"""
    fLat: float
    fLon: float  
    fHeight: float
    fSpeed: float
    roll: float  # 翻滚角
    pitch: float # 仰俯角
    ID: int
    timestamp: float

@dataclass
class RadarInfo:
    """敌方雷达信息"""
    fLat: float
    fLon: float
    fHeight: float
    timestamp: float
    ID: int

class HighPrecisionThreatAssessment:
    """高精度威胁评估系统"""
    
    def __init__(self, max_history_size=50):
        # 状态历史存储
        self.target_history: Dict[int, deque] = {}
        self.ownship_history: deque = deque(maxlen=max_history_size)
        self.max_history_size = max_history_size
        
        # 卡尔曼滤波器参数
        self.state_dim = 12  # [x, y, z, vx, vy, vz, ax, ay, az, jx, jy, jz]
        self.obs_dim = 3     # 观测位置
        
        # 初始化滤波器参数
        self._initialize_kalman_parameters()
        
        # 威胁评估参数
        self.threat_params = {
            'base_distance_weight': 0.4,
            'base_speed_weight': 0.3,
            'base_angle_weight': 0.3,
            'max_threat_distance': 100000,  # 100km
            'critical_distance': 10000,     # 10km
            'max_closure_speed': 1000,      # 1000m/s
        }
    
    def _initialize_kalman_parameters(self):
        """初始化卡尔曼滤波器参数"""
        # 状态转移矩阵（将在运行时根据Δt动态计算）
        self.F_base = np.eye(self.state_dim)
        
        # 观测矩阵
        self.H = np.zeros((self.obs_dim, self.state_dim))
        self.H[0, 0] = 1  # x
        self.H[1, 1] = 1  # y  
        self.H[2, 2] = 1  # z
        
        # 过程噪声协方差
        self.Q = np.eye(self.state_dim) * 0.1
        
        # 观测噪声协方差
        self.R = np.eye(self.obs_dim) * 10
        
        # 状态协方差
        self.P = np.eye(self.state_dim) * 100
        
        # 目标状态存储
        self.target_states: Dict[int, np.ndarray] = {}
        self.target_covariances: Dict[int, np.ndarray] = {}

    def llh_to_ecef(self, lat: float, lon: float, height: float) -> np.ndarray:
        """LLH坐标转ECEF坐标（优化版本）"""
        lat_rad = np.radians(lat)
        lon_rad = np.radians(lon)
        
        # WGS84参数
        a = 6378137.0  # 半长轴
        e2 = 6.69437999014e-3  # 第一偏心率平方
        
        # 计算卯酉圈曲率半径
        N = a / np.sqrt(1 - e2 * np.sin(lat_rad)**2)
        
        # ECEF坐标计算
        x = (N + height) * np.cos(lat_rad) * np.cos(lon_rad)
        y = (N + height) * np.cos(lat_rad) * np.sin(lon_rad) 
        z = (N * (1 - e2) + height) * np.sin(lat_rad)
        
        return np.array([x, y, z])


    def ecef_to_enu(self, ecef_target: np.ndarray, ecef_origin: np.ndarray, 
                   lat_origin: float, lon_origin: float) -> np.ndarray:
        """ECEF坐标转ENU坐标系（局部东北天）"""
        # 计算相对位置
        dx = ecef_target - ecef_origin
        
        lat_rad = np.radians(lat_origin)
        lon_rad = np.radians(lon_origin)
        
        # 旋转矩阵
        R = np.array([
            [-np.sin(lon_rad), np.cos(lon_rad), 0],
            [-np.sin(lat_rad)*np.cos(lon_rad), -np.sin(lat_rad)*np.sin(lon_rad), np.cos(lat_rad)],
            [np.cos(lat_rad)*np.cos(lon_rad), np.cos(lat_rad)*np.sin(lon_rad), np.sin(lat_rad)]
        ])
        
        return R @ dx
    
    @staticmethod
    @numba.jit(nopython=True, fastmath=True)
    def compute_dynamic_F(dt: float) -> np.ndarray:
        """计算动态状态转移矩阵（使用numba加速）"""
        F = np.eye(12)
        
        # 位置-速度关系
        F[0, 3] = dt
        F[1, 4] = dt
        F[2, 5] = dt
        
        # 位置-加速度关系  
        F[0, 6] = 0.5 * dt**2
        F[1, 7] = 0.5 * dt**2
        F[2, 8] = 0.5 * dt**2
        
        # 位置-加加速度关系
        F[0, 9] = (1/6) * dt**3
        F[1, 10] = (1/6) * dt**3
        F[2, 11] = (1/6) * dt**3
        
        # 速度-加速度关系
        F[3, 6] = dt
        F[4, 7] = dt
        F[5, 8] = dt
        
        # 速度-加加速度关系
        F[3, 9] = 0.5 * dt**2
        F[4, 10] = 0.5 * dt**2
        F[5, 11] = 0.5 * dt**2
        
        # 加速度-加加速度关系
        F[6, 9] = dt
        F[7, 10] = dt
        F[8, 11] = dt
        
        return F
    
    @staticmethod
    @numba.jit(nopython=True, fastmath=True)  
    def compute_dynamic_Q(dt: float) -> np.ndarray:
        """计算动态过程噪声协方差（使用numba加速）"""
        Q = np.eye(12) * 0.1
        
        # 为高阶项增加噪声
        Q[6:9, 6:9] *= dt**2  # 加速度噪声
        Q[9:12, 9:12] *= dt**4  # 加加速度噪声
        
        return Q
    
    def update_kalman_filter(self, target_id: int, observation: np.ndarray, dt: float):
        """更新卡尔曼滤波器状态"""
        if target_id not in self.target_states:
            # 初始化状态
            self.target_states[target_id] = np.zeros(self.state_dim)
            self.target_states[target_id][:3] = observation
            self.target_covariances[target_id] = np.eye(self.state_dim) * 100
            return
        
        # 获取当前状态
        x = self.target_states[target_id]
        P = self.target_covariances[target_id]
        
        # 动态计算状态转移矩阵和过程噪声
        F = self.compute_dynamic_F(dt)
        Q = self.compute_dynamic_Q(dt)
        
        # 预测步骤
        x_pred = F @ x
        P_pred = F @ P @ F.T + Q
        
        # 更新步骤
        y = observation - self.H @ x_pred  # 新息
        S = self.H @ P_pred @ self.H.T + self.R  # 新息协方差
        K = P_pred @ self.H.T @ np.linalg.inv(S)  # 卡尔曼增益
        
        # 状态更新
        x_new = x_pred + K @ y
        P_new = (np.eye(self.state_dim) - K @ self.H) @ P_pred
        
        # 存储更新后的状态
        self.target_states[target_id] = x_new
        self.target_covariances[target_id] = P_new
    
    def calculate_high_precision_closure_speed(self, target_id: int, 
                                             current_distance: float, dt: float) -> float:
        """计算高精度接近速度"""
        if target_id not in self.target_history:
            return 0.0
        
        history = self.target_history[target_id]
        if len(history) < 2:
            return 0.0
        
        # 获取历史距离
        prev_distance = history[-2]['distance']
        
        # 基础接近速度
        basic_closure = (prev_distance - current_distance) / dt
        
        # 如果历史数据足够，计算加速度校正
        if len(history) >= 3:
            prev_prev_distance = history[-3]['distance']
            prev_closure = (prev_prev_distance - prev_distance) / dt
            current_closure = basic_closure
            
            # 相对加速度
            rel_acceleration = (current_closure - prev_closure) / dt
            
            # 校正项
            correction = 0.5 * rel_acceleration * dt
            return basic_closure + correction
        
        return basic_closure
    
    def calculate_dynamic_weights(self, distance: float, closure_speed: float, 
                                dt: float, history: list) -> Tuple[float, float, float]:
        """计算动态权重"""
        # 基础权重
        w_d = self.threat_params['base_distance_weight']
        w_v = self.threat_params['base_speed_weight']
        w_a = self.threat_params['base_angle_weight']
        
        if len(history) < 2:
            return w_d, w_v, w_a
        
        # 距离变化率影响
        prev_distance = history[-2]['distance']
        distance_rate = abs((distance - prev_distance) / dt) if dt > 0 else 0
        w_d *= (1 + 0.1 * distance_rate / 100)  # 距离变化越快，距离权重越高
        
        # 速度变化率影响
        if len(history) >= 3:
            prev_closure = history[-2]['closure_speed']
            current_closure = closure_speed
            speed_rate = abs((current_closure - prev_closure) / dt) if dt > 0 else 0
            w_v *= (1 + 0.2 * speed_rate / 50)  # 速度变化越快，速度权重越高
        
        # 归一化权重
        total = w_d + w_v + w_a
        return w_d/total, w_v/total, w_a/total
    
    def calculate_threat_trend(self, target_id: int, current_threat: float, dt: float) -> float:
        """计算威胁趋势"""
        if target_id not in self.target_history:
            return current_threat
        
        history = self.target_history[target_id]
        if len(history) < 2:
            return current_threat
        
        # 威胁变化率
        prev_threat = history[-2]['threat']
        threat_rate = (current_threat - prev_threat) / dt if dt > 0 else 0
        
        # 威胁加速度
        threat_acceleration = 0
        if len(history) >= 3:
            prev_prev_threat = history[-3]['threat']
            prev_rate = (prev_threat - prev_prev_threat) / dt
            threat_acceleration = (threat_rate - prev_rate) / dt if dt > 0 else 0
        
        # 趋势增强威胁评估
        prediction_time = 5.0  # 预测5秒后的威胁
        trend_enhanced = current_threat + threat_rate * prediction_time + 0.5 * threat_acceleration * prediction_time**2
        
        return max(0, min(100, trend_enhanced))
    
    def calculate_relative_angle_threat(self, target_enu: np.ndarray, 
                                      ownship_attitude: Tuple[float, float]) -> float:
        """计算相对角度威胁"""
        # 计算目标相对于我方的方位角和高低角
        range_xy = np.sqrt(target_enu[0]**2 + target_enu[1]**2)
        azimuth = np.degrees(np.arctan2(target_enu[0], target_enu[1]))  # 方位角（东偏北）
        elevation = np.degrees(np.arctan2(target_enu[2], range_xy))     # 高低角
        
        # 归一化角度威胁（前方0°威胁最大，后方180°威胁最小）
        azimuth_threat = max(0, 1 - abs(azimuth) / 180)
        elevation_threat = max(0, 1 - abs(elevation) / 90)
        
        return 0.6 * azimuth_threat + 0.4 * elevation_threat
    
    def assess_threat(self, radar_info: RadarInfo, ownship_info: PlatParaInfo) -> float:
        """评估单个雷达威胁系数"""
        target_id = radar_info.ID
        
        # 计算时间间隔
        current_time = radar_info.timestamp
        dt = 0.1  # 默认时间间隔
        
        if target_id in self.target_history and self.target_history[target_id]:
            last_time = self.target_history[target_id][-1]['timestamp']
            dt = max(0.001, current_time - last_time)  # 避免除零
        
        # 坐标转换
        ecef_ownship = self.llh_to_ecef(ownship_info.fLat, ownship_info.fLon, ownship_info.fHeight)
        ecef_target = self.llh_to_ecef(radar_info.fLat, radar_info.fLon, radar_info.fHeight)
        
        enu_target = self.ecef_to_enu(ecef_target, ecef_ownship, 
                                    ownship_info.fLat, ownship_info.fLon)
        
        # 计算距离
        distance = np.linalg.norm(enu_target)
        
        # 更新卡尔曼滤波器
        self.update_kalman_filter(target_id, enu_target, dt)
        
        # 获取滤波后的状态
        if target_id in self.target_states:
            filtered_state = self.target_states[target_id]
            filtered_velocity = np.linalg.norm(filtered_state[3:6])
        else:
            filtered_velocity = 0
        
        # 计算高精度接近速度
        closure_speed = self.calculate_high_precision_closure_speed(target_id, distance, dt)
        
        # 初始化历史记录
        if target_id not in self.target_history:
            self.target_history[target_id] = deque(maxlen=self.max_history_size)
        
        # 计算各维度威胁（归一化到0-1）
        # 距离威胁（越近威胁越大）
        distance_threat = max(0, 1 - distance / self.threat_params['max_threat_distance'])
        
        # 速度威胁（接近速度越大威胁越大）
        speed_threat = min(1, abs(closure_speed) / self.threat_params['max_closure_speed'])
        
        # 角度威胁
        angle_threat = self.calculate_relative_angle_threat(enu_target, 
                                                          (ownship_info.roll, ownship_info.pitch))
        
        # 计算动态权重
        w_d, w_v, w_a = self.calculate_dynamic_weights(distance, closure_speed, dt, 
                                                     list(self.target_history[target_id]))
        
        # 基础威胁计算
        base_threat = w_d * distance_threat + w_v * speed_threat + w_a * angle_threat
        
        # 应用趋势增强
        current_threat = base_threat * 100  # 转换为0-100的分数
        trend_enhanced_threat = self.calculate_threat_trend(target_id, current_threat, dt)
        
        # 存储当前状态
        current_state = {
            'timestamp': current_time,
            'distance': distance,
            'closure_speed': closure_speed,
            'threat': trend_enhanced_threat,
            'enu_position': enu_target
        }
        self.target_history[target_id].append(current_state)
        
        return trend_enhanced_threat
    
    def batch_assess_threats(self, radar_infos: List[RadarInfo], 
                           ownship_info: PlatParaInfo) -> Dict[int, float]:
        """批量评估威胁（优化性能）"""
        threats = {}
        
        # 预处理我方位置
        ecef_ownship = self.llh_to_ecef(ownship_info.fLat, ownship_info.fLon, ownship_info.fHeight)
        
        for radar_info in radar_infos:
            threat_level = self.assess_threat(radar_info, ownship_info)
            threats[radar_info.ID] = threat_level
        
        return threats
    
    def get_priority_targets(self, radar_infos: List[RadarInfo], 
                           ownship_info: PlatParaInfo, top_k: int = 5) -> List[Tuple[int, float]]:
        """获取优先级目标列表"""
        threats = self.batch_assess_threats(radar_infos, ownship_info)
        
        # 按威胁系数排序
        sorted_threats = sorted(threats.items(), key=lambda x: x[1], reverse=True)
        
        return sorted_threats[:top_k]

# 使用示例和性能测试
def demo():
    """演示如何使用威胁评估系统"""
    
    # 初始化系统
    threat_system = HighPrecisionThreatAssessment()
    
    # 模拟数据
    ownship = PlatParaInfo(
        fLat=39.9, fLon=116.4, fHeight=1000, fSpeed=200,
        roll=0, pitch=0, ID=0, timestamp=time.time()
    )
    
    radar_threats = []
    for i in range(10):
        # 模拟不同距离和速度的雷达威胁
        lat = 39.9 + np.random.uniform(-0.5, 0.5)
        lon = 116.4 + np.random.uniform(-0.5, 0.5)
        height = 1000 + np.random.uniform(-200, 200)
        
        radar = RadarInfo(
            fLat=lat, fLon=lon, fHeight=height,
            timestamp=time.time() + i * 0.1, ID=i
        )
        radar_threats.append(radar)
    
    # 性能测试
    start_time = time.time()
    
    # 单次评估
    for radar in radar_threats:
        threat_level = threat_system.assess_threat(radar, ownship)
        print(f"雷达 {radar.ID}: 威胁系数 = {threat_level:.2f}")
    
    # 批量评估
    threats_dict = threat_system.batch_assess_threats(radar_threats, ownship)
    
    # 获取优先级目标
    priority_targets = threat_system.get_priority_targets(radar_threats, ownship, top_k=3)
    print("\n优先级目标:")
    for target_id, threat in priority_targets:
        print(f"目标 {target_id}: 威胁系数 = {threat:.2f}")
    
    end_time = time.time()
    print(f"\n处理时间: {(end_time - start_time)*1000:.2f} ms")

if __name__ == "__main__":
    demo()
